import backtrader as bt
from .MyTT import MA, EMA, RSI, ATR,MACD,TAQ  # 导入需要的指标函数
import numpy as np

class MyTT_Indicator(bt.Indicator):
    """
    MyTT指标适配器的基类
    子类需要实现:
    - mytt_func: mytt库中的指标函数
    - lines: 定义输出线名称
    - params: 定义参数
    """
    def __init__(self):
        self.history = []  # 存储历史数据
        self.data_fields = [self.data.close]  # 默认只使用收盘价
        
    def prenext(self):
        """在数据不足时收集数据"""
        self.collect_data()
        
    def next(self):
        """主计算逻辑"""
        self.collect_data()
        
        if len(self.history) >= max(self._get_required_bars()):
            try:
                result = self.mytt_func(self.history, *self._get_mytt_params())
                self._update_lines(result)
            except Exception as e:
                print(f"指标计算错误: {e}")
                self._set_nan()
        else:
            self._set_nan()
    
    def collect_data(self):
        """收集所需的历史数据"""
        current_values = [d[0] for d in self.data_fields]
        self.history.append(current_values)
        
        if len(self.history) > max(self._get_required_bars()):
            self.history.pop(0)
    
    def _get_mytt_params(self):
        """获取传递给mytt函数的参数"""
        return [getattr(self.p, param) for param in self.params]
    
    def _get_required_bars(self):
        """获取需要的bar数"""
        required_bars = []
        for param in self.params:
            param_name, default_value = param
            if hasattr(self.p, param_name):
                required_bars.append(getattr(self.p, param_name))
        return required_bars
    
    def _update_lines(self, result):
        """
        更新指标线（关键修正点）
        处理mytt返回的numpy数组/标量/列表等情况
        """
        if isinstance(result, (np.ndarray, list, tuple)):
            for i, line in enumerate(self.lines):
                if i < len(result):
                    line[0] = float(result[i]) if len(result) > i else float('nan')
        else:
            # 如果是标量值，更新第一条线
            if len(self.lines) > 0:
                self.lines[0][0] = float(result)
    
    def _set_nan(self):
        """数据不足时设置为NaN"""
        for line in self.lines:
            line[0] = float('nan')
 


# 实际的子类
class MyTT_MA(MyTT_Indicator):
    lines = ('ma',)
    params = (('period', 5),)
    
    def mytt_func(self, data, period):
        close_prices = [d[0] for d in data]  # 提取收盘价
        return MA(close_prices, period)

class MyTT_EMA(MyTT_Indicator):
    lines = ('ema',)
    params = (('period', 5),)
    
    def mytt_func(self, data, period):
        close_prices = [d[0] for d in data]
        return EMA(close_prices, period)

class MyTT_MACD(MyTT_Indicator):
    lines = ('dif', 'dea','macd')
    params = (('short', 12), ('long', 26), ('signal', 9))
    
    def mytt_func(self, data, short, long, signal):
        close_prices = [d[0] for d in data]  # 提取收盘价
        result = MACD(close_prices, short, long, signal)
        print(f"Macd 返回值{result}")
        return result
        

class MyTT_RSI(MyTT_Indicator):
    lines = ('rsi',)
    params = (('period', 14),)
    
    def mytt_func(self, data, period):
        close_prices = [d[0] for d in data]
        return RSI(close_prices, period)

class MyTT_ATR(MyTT_Indicator):
    lines = ('atr',)
    params = (('period', 14),)
    
    def __init__(self):
        self.data_fields = [self.data.high, self.data.low, self.data.close]
        super().__init__()
    
    def mytt_func(self, data, period):
        highs = [d[0] for d in data]
        lows = [d[1] for d in data]
        closes = [d[2] for d in data]
        return ATR(highs, lows, closes, period)
    
